import asyncio
from typing import Optional, final
from contextlib import AsyncExitStack

from services.agent_service import  create_llm, create_agent, create_llm2

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from langchain.tools import Tool, StructuredTool
from langgraph.prebuilt import create_react_agent
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain.agents import create_structured_chat_agent
from langchain import hub
from langchain_ollama import OllamaLLM, ChatOllama
from langgraph.graph import MessageGraph
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI




from dotenv import load_dotenv

load_dotenv()  # load environment variables from .env

class MCPClient:
    def __init__(self):
        # Initialize session and client objects
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()
        # self.anthropic = Anthropic()
        self.llm = create_llm2()

    async def connect_to_server(self, server_script_path: str):
        """Connect to an MCP server
        
        Args:
            server_script_path: Path to the server script (.py or .js)
        """
        is_python = server_script_path.endswith('.py')
        is_js = server_script_path.endswith('.js')
        if not (is_python or is_js):
            raise ValueError("Server script must be a .py or .js file")
            
        command = "python3" if is_python else "node"
        server_params = StdioServerParameters(
            command=command,
            args=[server_script_path],
            env=None
        )
        
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
        
        await self.session.initialize()
        self.tools = await load_mcp_tools(self.session)
        # print("self.tools:", self.tools)
        # List available tools
        response = await self.session.list_tools()
        tools = response.tools
        print("\nConnected to server with tools:", [tool.name for tool in tools])

    async def process_query(self, query: str) -> str:
        """使用langgraph的ReAct代理和本地Ollama模型处理查询"""


        response = await self.llm.bind_tools(self.tools).ainvoke(query)
        # print("response step 1:", response)

        # 现在有一个会计凭证编号为1234，帮我生成摘要并返回文件路径；用中文回答
        # 现在有一个会计凭证编号为abcd，帮我生成摘要并返回文件路径；用中文回答

        # Process response and handle tool calls
        tool_results = []
        final_text = []


        # print("type",  content.type)
        if hasattr(response, 'tool_calls'):
            # print("tool_calls", response.tool_calls)

            for tool_call in response.tool_calls:
                # print("tool_call", tool_call)
                # print("type",  type(tool_call))
                tool_name = tool_call['name'] 
                tool_args = tool_call['args']
                # print("tool_args",  tool_args)
                # print("tool_args type",  type(tool_args))
                # print("accdoccode",  tool_args['accdoccode'])
                # print("accdoccode type",  type(tool_args['accdoccode']))
            
                # Execute tool call
                tool_result = await self.session.call_tool(tool_name, tool_args)
                # print(f"call_tool {tool_name} result", tool_result)
                tool_results.append({"call": tool_name, "result": tool_result})
                # print("tool_result isError:", tool_result.isError)
                # print("tool_result content", tool_result.content)
                # print("tool_result content type", type(tool_result.content))
                text_str = '\n'.join(str(item.text) for item in tool_result.content)
                print(f"call_tool {tool_name} text_str:", text_str)
                final_text.append(f"Calling tool {tool_name} with args {tool_args}, get result: {text_str}")


                    # Continue conversation with tool results
                    # if hasattr(content, 'text') and content.text:
                    #     messages.append({
                    #     "role": "assistant",
                    #     "content": content.text
                    #     })
                    # messages.append({
                    #     "role": "user", 
                    #     "content": result.content
                    # })

                    # # Get next response from Claude
                    # response = self.anthropic.messages.create(
                    #     model="claude-3-5-sonnet-20241022",
                    #     max_tokens=1000,
                    #     messages=messages,
                    # )

                    # final_text.append(response.content[0].text)
        
        
            final_message = f"用户问题：{query}, 已知提示：" + "\n".join(final_text)
            response2 = await self.llm.bind_tools(self.tools).ainvoke(final_message)
            # print("response2", response2)
            return response2.content

        else:
            print("No tool_calls in response")
            return response.content


    async def chat_loop(self):
        """Run an interactive chat loop"""
        print("\nMCP Client Started!")
        print("Type your queries or 'quit' to exit.")
        
        while True:
            try:
                query = input("\nQuery: ").strip()
                
                if query.lower() == 'quit':
                    break
                    
                response = await self.process_query(query)
                print("\n" + response)
                    
            except Exception as e:
                print(f"\nError: {str(e)}")
    
    async def cleanup(self):
        """Clean up resources"""
        await self.exit_stack.aclose()

async def main():
    if len(sys.argv) < 2:
        print("Usage: python client.py <path_to_server_script>")
        sys.exit(1)
        
    client = MCPClient()
    try:
        await client.connect_to_server(sys.argv[1])
        await client.chat_loop()
    finally:
        await client.cleanup()

if __name__ == "__main__":
    import sys
    asyncio.run(main())